#include "yolo_v6.h"
#include "../txr_algo_dlm_v3/src/common/common.hpp"
#include "yaml-cpp/yaml.h"

#include "boost/process.hpp"
#include "boost/process/windows.hpp"
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/path.hpp"
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/fstream.hpp"

#include "dll.h"

using namespace CryptoPP;

std::vector<char> decryptFile(const std::vector<char> & v_encrypt, const std::string& key, const std::string& iv)
{
	std::vector<char> v_decrypt;
	try
	{
		CBC_Mode<AES>::Decryption decryption((const byte*)key.data(), key.size(), (const byte*)iv.data());

		std::vector<char> _decrypt = v_encrypt;
		decryption.ProcessData(
			reinterpret_cast<CryptoPP::byte*>(_decrypt.data()),
			reinterpret_cast<const CryptoPP::byte*>(v_encrypt.data()),
			v_encrypt.size());
		
		byte paddingByte = _decrypt.back();
		size_t paddingSize = static_cast<size_t>(paddingByte);
		_decrypt.resize(v_encrypt.size() - paddingSize);

		v_decrypt = _decrypt;
	}
	catch (const CryptoPP::Exception& e)
	{
		std::cerr << e.what() << std::endl;
	}
	return v_decrypt;
}


yolo_v6::yolo_v6()
{

}
yolo_v6::~yolo_v6()
{

}

std::vector<char> get_file_data(std::string file_path)
{
	std::ifstream file(file_path, std::ios_base::binary);
	assert(file.good());
	file.seekg(0, std::ios_base::end);
	int file_len = file.tellg();
	file.seekg(0, std::ios_base::beg);

	std::vector<char> data(file_len, 0);
	file.read(data.data(), file_len);

	file.close();
	return data;
}

void write_file_data(std::vector<char> data, std::string file_path)
{
	std::ofstream file(file_path, std::ios_base::binary);
	file.write(data.data(), data.size());
	file.close();
}

void encrypt_data(std::vector<char> data)
{
	char trans = 55;
	for (int i = 0; i < data.size(); ++i)
	{
		data[i] += trans;
		trans += 11;
	}
}

void encryptFile(st_encrypt_info info, std::vector<char> engine_data, std::string target_path)
{
	info.engine_size = engine_data.size();

	std::vector<char> target_data(sizeof(info), 0);
	*((st_encrypt_info*)(target_data.data())) = info;

	target_data.insert(target_data.end(), engine_data.begin(), engine_data.end());

	encrypt_data(target_data);
	write_file_data(target_data, target_path);
	std::cout << "writing encrypt file..." << target_path <<" done!" << std::endl;

}


bool yolo_v6::TransModel(st_trans_model_input input)
{
	namespace fs = boost::filesystem;
	YAML::Node root = YAML::LoadFile(input.cfg_path);
	YAML::Node config;
	if (root["YOLOv6"].IsDefined())
	{
		config = root["YOLOv6"];
	}
	if (root["DET_LB"].IsDefined())
	{
		config = root["DET_LB"];
	}
	//std::string onnx_file = config["onnx_file"].as<std::string>();
	//std::string engine_file = config["engine_file"].as<std::string>();
	int batch_size = config["BATCH_SIZE"].as<int>();
	if (!config["CLASS_NUM"].IsDefined())
	{
		std::cout << "CLASS_NUM undefined" << std::endl;
	}
	int class_num = config["CLASS_NUM"].as<int>();
	int input_channel = config["INPUT_CHANNEL"].as<int>();
	int image_width = config["IMAGE_WIDTH"].as<int>();
	int image_height = config["IMAGE_HEIGHT"].as<int>();
	float obj_threshold = config["obj_threshold"].as<float>();
	float nms_threshold = config["nms_threshold"].as<float>();
	std::vector<int> strides = config["strides"].as<std::vector<int>>();
	
	std::vector<int> num_anchors(strides.size(),1);
	if (config["num_anchors"].IsDefined())
	{
		std::vector<int> anchors = config["num_anchors"].as<std::vector<int>>();
		if (anchors.size() == strides.size())
		{
			num_anchors = anchors;
		}
	}
	std::vector<std::vector<int>> grids;
	for (int index = 0; index < strides.size(); ++index)
	{
		grids.push_back({ int(image_height / strides[index]), int(image_width / strides[index])* num_anchors[index] });
	}

	std::vector<char> v_model, v_engine;

	fs::path p(input.model_path.c_str());
	if (p.extension().string() == ".onnx")
	{
		v_model = get_file_data(p.string());
	}
	else if (p.extension().string() == ".ennx")
	{
		std::vector<char> v_e_model;
		v_e_model = get_file_data(p.string());
		std::string key = "irkdh_573%3?iq5h";
		std::string iv = "1234567890123456";
		v_model = decryptFile(v_e_model, key, iv);
	}
	else
	{
		return false;
	}
	{
		nvinfer1::ICudaEngine* engine = nullptr;
		onnxToTRTModel(v_model, v_engine, engine, batch_size);
		assert(engine != nullptr);
		if (engine == nullptr)
		{
			return false;
		}
		engine->destroy();
	}

	st_encrypt_info info;
	sprintf_s(info.gpu_model, 100, "%s", input.gpu_model.c_str());
	info.batch_size = batch_size;
	info.class_num = class_num;
	info.input_channel = input_channel;
	info.image_width = image_width;
	info.image_height = image_height;
	info.stride_num = strides.size();
	for (int i = 0;i<strides.size();++i)
	{
		info.strides[i] = strides[i];
		info.grid[i][0] = grids[i][0];
		info.grid[i][1] = grids[i][1];
	}
	std::string encrypt_file = input.trans_path;
	encryptFile(info, v_engine, encrypt_file);

	return true;
}